import cv2
import os
import sys
import imageio

import Utils
import Segmentation
import Cell
import Tarcking


def main(pic_path, gen_path):
    # this is for store all of the image data
    img_draw_list = []
    # this stores cells information of all of the image
    all_cells_list = []

    count = 1
    for filename in sorted(os.listdir(pic_path)):
        # img is used to store the image data
        img_color = cv2.imread(pic_path + "/" + filename)
        img_gray = cv2.cvtColor(img_color, cv2.COLOR_BGR2GRAY)

        # pre_process
        image_pre_process, markers = Segmentation.pre_process(img_color)
        # image_pre_process = Segmentation.pre_process_th(img_gray)

        # calculate contour
        contours = Segmentation.calc_contour(image_pre_process, markers)
        # img_contour, contours, hierarchy = Segmentation.calc_contours_(image_pre_process)

        # create cell object
        color_list = Utils.get_color_list(len(contours))
        cell_list = Cell.save_cells_info(contours, color_list)

        # match and register id
        cell_id = 0
        if len(all_cells_list) < 1:  # the first frame
            for cell in cell_list:
                cell.register_id(cell_id)
                cell_id = cell_id + 1
        else:
            cell_list = Tarcking.cell_match(all_cells_list[-1], cell_list)

        # all_cells_list save all cell of a image
        all_cells_list.append(cell_list)

        # draw and save image
        img_bgr = Segmentation.draw_cell_contour_and_ID(img_gray, cell_list)
        img_bgr = Tarcking.draw_path(img_bgr, cell_list)
        img_bgr = Utils.draw_img_info(img_bgr, cell_list)
        cv2.imwrite(f'{gen_path}/{count}.tif', img_bgr)

        count = count + 1

        # BGR to RGB
        img_draw_list.append(img_bgr[..., ::-1])

    # show img as RGB
    gif = imageio.mimsave('./result.gif', img_draw_list, 'GIF', duration=0.01)


if __name__ == '__main__':

    # pic_path = "F:/UNSW/9517/group/Sequences/01"
    # gen_path = "F:/UNSW/9517/group/Sequences/01_gen"
    # if len(sys.argv) > 2:
    #     pic_path = sys.argv[1]
    #     gen_path = sys.argv[2]
    #
    # if os.path.isdir(pic_path):
    #     print("Now the gen path is ", gen_path)
    #     if not os.path.exists(gen_path):
    #         os.makedirs(gen_path)
    #         print("gen path created successfully")
    #     main(pic_path, gen_path)
    # else:
    #     print("pic_path is not found!")

    for i in range(0, 5):
        pic_path = "F:/UNSW/9517/group/Sequences/0" + str(i)
        gen_path = "F:/UNSW/9517/group/Sequences/0" + str(i) + "_gen"
        if len(sys.argv) > 2:
            pic_path = sys.argv[1]
            gen_path = sys.argv[2]

        if os.path.isdir(pic_path):
            print("Now the gen path is ", gen_path)
            if not os.path.exists(gen_path):
                os.makedirs(gen_path)
                print("gen path created successfully")
            main(pic_path, gen_path)
        else:
            print("pic_path is not found!")
